# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import numpy as np
from numpy.random import choice
from scipy import interpolate
import random
import cv2
from skfmm import distance
from isaacgym import terrain_utils
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg
from scipy import ndimage
from isaacgym import gymutil, gymapi
from legged_gym.utils import trimesh
# from pydelatin import Delatin
# import pyfqmr
from scipy.ndimage import binary_dilation
from matplotlib.path import Path
from legged_gym.utils.terrain_modules import *
import legged_gym.utils.terrain_nav_modules as nav_modules

def rotate_point(point, angle, origin):
    """Rotate a point by a given angle around an origin."""
    point = np.array(point) - origin
    rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)], 
                                [np.sin(angle),  np.cos(angle)]])
    rotated_point = np.dot(rotation_matrix, point)
    return rotated_point + origin

def generate_forbidden_edge(horizontal_scale, vertical_scale, num_rows, num_cols, hf, slope_threshold):
    limit_threshold = slope_threshold
    limit_threshold *= horizontal_scale / vertical_scale
    forbidden_x = np.zeros((num_rows, num_cols))
    forbidden_y = np.zeros((num_rows, num_cols))
    forbidden_corners = np.zeros((num_rows, num_cols))
    forbidden_x[:num_rows-1, :] += (hf[1:num_rows, :] - hf[:num_rows-1, :] > limit_threshold)
    forbidden_x[1:num_rows, :] -= (hf[:num_rows-1, :] - hf[1:num_rows, :] > limit_threshold)
    forbidden_x = (forbidden_x != 0)
    forbidden_y[:, :num_cols-1] += (hf[:, 1:num_cols] - hf[:, :num_cols-1] > limit_threshold)
    forbidden_y[:, 1:num_cols] -= (hf[:, :num_cols-1] - hf[:, 1:num_cols] > limit_threshold)
    forbidden_y = (forbidden_y != 0)
    forbidden_corners[:num_rows-1, :num_cols-1] += (hf[1:num_rows, 1:num_cols] - hf[:num_rows-1, :num_cols-1] > limit_threshold)
    forbidden_corners[1:num_rows, 1:num_cols] -= (hf[:num_rows-1, :num_cols-1] - hf[1:num_rows, 1:num_cols] > limit_threshold)
    forbidden_corners = (forbidden_corners != 0)

    forbidden_area = forbidden_x | forbidden_y | forbidden_corners
    
    return forbidden_area

def fill_rectangles(arr):
    # 寻找轮廓
    contours, _ = cv2.findContours(arr.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 创建与输入数组相同大小的填充数组
    filled_array = np.zeros_like(arr, dtype=np.uint8)

    # 填充轮廓内部
    cv2.fillPoly(filled_array, contours, color=(255, 255, 255))
    
    filled_array_bool = filled_array.astype(bool)
    
    return filled_array_bool

class Terrain:
    def __init__(self, cfg: LeggedRobotCfg.terrain, num_robots, sim=None, gym=None) -> None:
        self.cfg = cfg
        self.num_robots = num_robots
        self.type = cfg.mesh_type
        if self.type in ["none", 'plane']:
            return
        self.env_length = cfg.terrain_length
        self.env_width = cfg.terrain_width
        
        self.nav_terrain_length = cfg.nav_terrain_length
        self.nav_terrain_width = cfg.nav_terrain_width
        self.virtual = self.cfg.virtual

        cfg.terrain_proportions = np.array(cfg.terrain_proportions) / np.sum(cfg.terrain_proportions)
        self.proportions = [np.sum(cfg.terrain_proportions[:i+1]) for i in range(len(cfg.terrain_proportions))]
        self.env_origins = np.zeros((cfg.num_rows, cfg.num_cols, cfg.n_cols_per_terrain, 3))
        self.env_final_goals = np.zeros((cfg.num_rows, cfg.num_cols, cfg.n_cols_per_terrain, 3))
        self.terrain_type = np.zeros((cfg.num_rows, cfg.num_cols))
        if self.cfg.waypoint_info_dim > 0:
            self.goals = np.zeros((cfg.num_rows, cfg.num_cols, cfg.num_goals, 3+self.cfg.waypoint_info_dim))
        else:
            self.goals = np.zeros((cfg.num_rows, cfg.num_cols, cfg.num_goals, 3))
        self.num_goals = cfg.num_goals

        self.width_per_env_pixels = int(self.env_width / cfg.horizontal_scale)
        self.length_per_env_pixels = int(self.env_length / cfg.horizontal_scale)
        self.nav_terrain_width_pixel = int(cfg.nav_terrain_width / cfg.horizontal_scale)  
        self.nav_terrain_length_pixel = int(cfg.nav_terrain_length / cfg.horizontal_scale)
        self.nav_block_length_pixel = int(cfg.block_length / cfg.horizontal_scale)
        self.nav_block_width_pixel = int(cfg.block_width / cfg.horizontal_scale)
        self.distances_fields = np.zeros((cfg.num_rows, cfg.num_cols, self.nav_terrain_length_pixel, self.nav_terrain_width_pixel))
        self.obstacle_maps = np.zeros((cfg.num_rows, cfg.num_cols, self.nav_terrain_length_pixel, self.nav_terrain_width_pixel))
        self.sample_reject_areas = np.zeros((cfg.num_rows, cfg.num_cols, self.nav_terrain_length_pixel, self.nav_terrain_width_pixel))
        self.terrain_offsets = np.zeros((cfg.num_rows, cfg.num_cols, 3))

        self.border = int(cfg.border_size/self.cfg.horizontal_scale)
        self.tot_cols = int(cfg.num_cols * self.nav_terrain_width_pixel) + 2 * self.border
        self.tot_rows = int(cfg.num_rows * self.nav_terrain_length_pixel) + 2 * self.border

        self.height_field_raw = np.zeros((self.tot_rows , self.tot_cols), dtype=np.int16)
        self.virtual_height_field = np.zeros((self.tot_rows , self.tot_cols), dtype=np.int16)  # obstacles
        if cfg.nav_terrain:
            self.gen_nav_terrain()
        else:    
            raise NotImplementedError("Only nav terrain is supported") 
        
        # post process
        # self.heightsamples = self.height_field_raw
        self.heightsamples = self.virtual_height_field.copy()
        # self.height_field_raw = self.virtual_height_field.copy()
        
        self.terrain_offsets *= self.cfg.horizontal_scale
        
        if self.type=="trimesh":
            print("Converting heightmap to trimesh...")
            if cfg.hf2mesh_method == "grid":
                self.vertices, self.triangles, self.edge_mask = convert_heightfield_to_trimesh(   self.height_field_raw,
                                                                                                self.cfg.horizontal_scale,
                                                                                                self.cfg.vertical_scale,
                                                                                                self.cfg.slope_treshold,)
                half_edge_width = int(self.cfg.edge_width_thresh / self.cfg.horizontal_scale)
                structure = np.ones((half_edge_width*2+1, 1))
                self.edge_mask = binary_dilation(self.edge_mask, structure=structure)
            else:
                assert cfg.hf2mesh_method == "fast", "Height field to mesh method must be grid or fast"
                self.vertices, self.triangles = convert_heightfield_to_trimesh_delatin(self.height_field_raw, self.cfg.horizontal_scale, self.cfg.vertical_scale, max_error=cfg.max_error)
            print("Created {} vertices".format(self.vertices.shape[0]))
            print("Created {} triangles".format(self.triangles.shape[0]))
    
    def gen_nav_terrain(self):
        for i in range(self.cfg.num_rows):
            for j in range(self.cfg.num_cols):
                if hasattr(self.cfg, "max_difficulty"):
                    difficulty = self.cfg.max_difficulty
                else:
                    difficulty = i / (self.cfg.num_rows)
                choice = j / self.cfg.num_cols + 0.001
                terrain = self.make_nav_terrain(difficulty, choice)
                self.add_terrain_to_map(terrain, i, j)
    
    def add_roughness(self, terrain, difficulty=0.2, virtual=False):
        max_height = (self.cfg.height[1] - self.cfg.height[0]) * difficulty + self.cfg.height[0]
        height = random.uniform(self.cfg.height[0], max_height)
        terrain_utils.random_uniform_terrain(terrain, min_height=-height, max_height=height, step=0.005, downsampled_scale=self.cfg.downsampled_scale, virtual=virtual)
    
    def make_nav_terrain(self, difficulty, choice=None):
        if random.random() < 0.5:
            self.options = self.cfg.nav_stair_options
        else:
            self.options = self.cfg.nav_step_options
        terrain = terrain_utils.SubTerrain("terrain",
                                width=self.nav_terrain_length_pixel,
                                length=self.nav_terrain_width_pixel,
                                vertical_scale=self.cfg.vertical_scale,
                                horizontal_scale=self.cfg.horizontal_scale)
        if getattr(self.cfg, "terrain_template", None) is not None and choice is not None:
            if choice < 0.5:
                nav_terrain(terrain, self.cfg, self.cfg.n_rows_per_terrain, self.cfg.n_cols_per_terrain, self.nav_block_length_pixel, self.nav_block_width_pixel, difficulty, self.options, self.virtual)
            else:
                nav_terrain(terrain, self.cfg, self.cfg.n_rows_per_terrain, self.cfg.n_cols_per_terrain, self.nav_block_length_pixel, self.nav_block_width_pixel, difficulty, self.options, self.virtual)
        else:
            nav_terrain(terrain, self.cfg, self.cfg.n_rows_per_terrain, self.cfg.n_cols_per_terrain, self.nav_block_length_pixel, self.nav_block_width_pixel, difficulty, self.options, self.virtual)
        idx = -1
        self.add_roughness(terrain, virtual=self.virtual)
        terrain.idx = idx
        return terrain

    def add_terrain_to_map(self, terrain, row, col):
        i = row
        j = col
        # map coordinate system
        start_x = self.border + i * self.nav_terrain_length_pixel
        end_x = self.border + (i + 1) * self.nav_terrain_length_pixel
        start_y = self.border + j * self.nav_terrain_width_pixel
        end_y = self.border + (j + 1) * self.nav_terrain_width_pixel
        self.height_field_raw[start_x: end_x, start_y:end_y] = terrain.height_field_raw
        self.virtual_height_field[start_x: end_x, start_y:end_y] = terrain.virtual_height_field
        self.terrain_offsets[i, j] = [start_x-self.border, start_y-self.border, 0]

        sample_width_thresh = int(self.cfg.sample_width_thresh / self.cfg.horizontal_scale)
        structure = np.ones((sample_width_thresh*2+1, sample_width_thresh*2+1))
        forbidden_edge = generate_forbidden_edge(terrain.horizontal_scale, terrain.vertical_scale, terrain.height_field_raw.shape[0], 
                                terrain.height_field_raw.shape[1], terrain.height_field_raw, self.cfg.sample_slope_treshold)
        obstacle_edge = generate_forbidden_edge(terrain.horizontal_scale, terrain.vertical_scale, terrain.height_field_raw.shape[0], 
                                terrain.height_field_raw.shape[1], terrain.height_field_raw, self.cfg.forbidden_slope_treshold)
        obstacle_area = fill_rectangles(obstacle_edge)

        if "forbidden" in self.options and getattr(self.cfg, "terrain_template", None) is None:
            sample_reject_area = forbidden_edge | obstacle_area
        else:
            sample_reject_area = forbidden_edge
        sample_reject_area = binary_dilation(sample_reject_area, structure=structure)
        sample_reject_area = binary_dilation(sample_reject_area, structure=structure)
        # sample_reject_area[:10, :] = 1; sample_reject_area[-10:, :] = 1
        # sample_reject_area[:, :10] = 1; sample_reject_area[:, -10:] = 1
        self.sample_reject_areas[i, j] = sample_reject_area

        # if "forbidden" in self.cfg.nav_options and getattr(self.cfg, "terrain_template", None) is None:
        #     obstacle_area = ~obstacle_area
        #     distances_field = distance(obstacle_area.astype(np.float32))
        #     self.distances_fields[i, j] = distances_field * self.cfg.horizontal_scale
        #     self.obstacle_maps[i, j] = ~obstacle_area
        #     self.obstacle_maps[i, j][:3, :] = 1
        #     self.obstacle_maps[i, j][-3:, :] = 1
        #     self.obstacle_maps[i, j][:, :3] = 1
        #     self.obstacle_maps[i, j][:, -3:] = 1

        # env_origin random generate in the first row
        env_origin_x = i * self.nav_terrain_length + 1.0
        # env_origin_y = (j + 0.5) * self.nav_terrain_width
        x1 = int((self.nav_terrain_length/2. - 0.5) / terrain.horizontal_scale) # within 1 meter square range
        x2 = int((self.nav_terrain_length/2. + 0.5) / terrain.horizontal_scale)
        y1 = int((self.nav_terrain_width/2. - 0.5) / terrain.horizontal_scale)
        y2 = int((self.nav_terrain_width/2. + 0.5) / terrain.horizontal_scale)
        if self.cfg.origin_zero_z:
            env_origin_z = 0
        else:
            env_origin_z = np.max(terrain.height_field_raw[x1:x2, y1:y2])*terrain.vertical_scale
        
        env_origin_y = np.array([j * self.nav_terrain_width + (item + 0.5) * self.cfg.block_width for item in range(self.cfg.n_cols_per_terrain)])
        if self.cfg.fixed_origins:
            tmp_y = j * self.nav_terrain_width + 0.5 * self.cfg.block_width
            origin_y = np.array([tmp_y]).repeat(self.cfg.n_cols_per_terrain)
        else:
            origin_y = env_origin_y
        origins = np.vstack((np.array([env_origin_x]).repeat(self.cfg.n_cols_per_terrain), 
                             origin_y, 
                             np.array([env_origin_z]).repeat(self.cfg.n_cols_per_terrain))).transpose(1,0)
        self.env_origins[i, j] = origins
        
        if self.cfg.fixed_final_goals:
            tmp_y = j * self.nav_terrain_width + 0.5 * self.cfg.block_width + self.cfg.block_width * 2
            goal_y = np.array([tmp_y]).repeat(self.cfg.n_cols_per_terrain)
        else:
            goal_y = env_origin_y
        env_goal_x = env_origin_x + self.nav_terrain_length - self.cfg.block_length * 3
        self.env_final_goals[i, j] = np.vstack((np.array([env_goal_x]).repeat(self.cfg.n_cols_per_terrain),
                                            goal_y,
                                            np.array([env_origin_z]).repeat(self.cfg.n_cols_per_terrain))).transpose(1,0)
                
        self.goals[i, j, :terrain.goals.shape[0], :2] = terrain.goals[:, :2] + [i * self.nav_terrain_length, j * self.nav_terrain_width]

def generate_obstacle_order(self, options, n_rows_per_track, n_cols_per_track, cfg):
    arr = np.random.choice(
            len(options), 
            size=(n_rows_per_track, n_cols_per_track))
    replacement_values = np.zeros(arr.shape[1])

    # not up & down & forbidden
    for i in range(arr.shape[0]):  # 遍历每一行
        if np.any(arr[i] == len(options)-2) and np.any(arr[i] == len(options)-1):  # 若上升和下降同时出现
            if np.random.choice([True, False]):
                arr[i] = np.where(arr[i] == len(options)-2, replacement_values, arr[i])
            else:
                arr[i] = np.where(arr[i] == len(options)-1, replacement_values, arr[i])  
    if getattr(cfg, "single_fine_tune", None) is not None:
        arr[1,:] = 1
    return arr

def nav_terrain(terrain, cfg, n_rows_per_terrain, n_cols_per_terrain, block_length, block_width, difficulty, options=None, virtual=False):
    if getattr(cfg, "terrain_template", None) is None:
        # import ipdb; ipdb.set_trace()
        obstacle_order = generate_obstacle_order(terrain, options, n_rows_per_terrain, n_cols_per_terrain, cfg)
        # print(obstacle_order)
    else:
        obstacle_order = np.array(cfg.terrain_template)
        cfg.n_rows_per_terrain = obstacle_order.shape[0]
        # goal_lists = cfg.goal_lists[final_goal]
        # print(goal_lists, final_goal)
    
    # buidl wall
    temp_goals = {}
    terrain_goals = np.zeros((1, 3))
    wall_thickness = np.random.uniform(*cfg.wall_thickness) if isinstance(cfg.wall_thickness, (tuple, list)) else cfg.wall_thickness
    wall_thickness = np.ceil(wall_thickness / cfg.horizontal_scale ).astype(int)
    wall_height = np.random.uniform(*cfg.wall_height) if isinstance(cfg.wall_height, (tuple, list)) else cfg.wall_height
    if cfg.if_wall:        
        terrain.height_field_raw[:wall_thickness, :] = wall_height / cfg.vertical_scale 
        terrain.height_field_raw[-wall_thickness:, :] = wall_height / cfg.vertical_scale 
        terrain.height_field_raw[:, :wall_thickness] = wall_height / cfg.vertical_scale 
        terrain.height_field_raw[:, -wall_thickness:] = wall_height / cfg.vertical_scale 
        if virtual:
            terrain.virtual_height_field[:wall_thickness, :] = wall_height / cfg.vertical_scale 
            terrain.virtual_height_field[-wall_thickness:, :] = wall_height / cfg.vertical_scale 
            terrain.virtual_height_field[:, :wall_thickness] = wall_height / cfg.vertical_scale 
            terrain.virtual_height_field[:, -wall_thickness:] = wall_height / cfg.vertical_scale        
    
    row_base_height = 0
    for local_row_id in range(cfg.n_rows_per_terrain+2):
        block_starting_height_px = row_base_height
        add_flag = True
        for local_col_id in range(cfg.n_cols_per_terrain):
            heightfield_x0 = block_length * local_row_id + wall_thickness
            heightfield_y0 = block_width * local_col_id + wall_thickness
            heightfield_x1 = block_length * (local_row_id + 1) + wall_thickness
            heightfield_y1 = block_width * (local_col_id + 1) + wall_thickness
            if local_row_id == 0 or local_row_id == cfg.n_rows_per_terrain + 1:
                track_heightfield, virtual_height_field, block_info, height_offset_px, goals = getattr(nav_modules, "get_" + "flat" + "_block")(
                cfg,
                difficulty = difficulty,
                heightfield_template = terrain.height_field_raw[
                    heightfield_x0: heightfield_x1,
                    heightfield_y0: heightfield_y1],
                block_length = block_length,
                block_width = block_width,
                block_starting_height_px = block_starting_height_px,
                )
                # if (local_row_id, local_col_id) in goal_lists and local_row_id == cfg.n_rows_per_terrain + 1:
                #     goals[:] += np.array([heightfield_x0*cfg.horizontal_scale, heightfield_y0*cfg.horizontal_scale, 0])
                #     terrain_goals = np.vstack((terrain_goals, goals))
                # if goals is not None:
                #     goals[:] += np.array([heightfield_x0*cfg.horizontal_scale, heightfield_y0*cfg.horizontal_scale, 0])
                #     temp_goals[(local_row_id, local_col_id)] = goals
            else:
                obstacle_selection = obstacle_order[local_row_id-1, local_col_id]
                obstacle_name = options[obstacle_selection]
                # obstacle_id = cfg.track_options_id_dict[obstacle_name]
                # if local_col_id != 0 and local_col_id != self.n_cols_per_track - 1:
                #     no_wall = True
                #     print("Warning: no wall for the track block.")
                track_heightfield, virtual_height_field, block_info, height_offset_px, goals = getattr(nav_modules, "get_" + obstacle_name + "_block")(
                    cfg,
                    difficulty = difficulty,
                    heightfield_template = terrain.height_field_raw[
                        heightfield_x0: heightfield_x1,
                        heightfield_y0: heightfield_y1],
                    block_length = block_length,
                    block_width = block_width,
                    block_starting_height_px = block_starting_height_px,
                )
                # if (local_row_id-1, local_col_id) in goal_lists:
                #     goals[:] += np.array([heightfield_x0*cfg.horizontal_scale, heightfield_y0*cfg.horizontal_scale, 0])
                #     terrain_goals = np.vstack((terrain_goals, goals))
                # if goals is not None:
                #     goals[:] += np.array([heightfield_x0*cfg.horizontal_scale, heightfield_y0*cfg.horizontal_scale, 0])
                #     temp_goals[(local_row_id, local_col_id)] = goals
                if height_offset_px != 0 and add_flag and obstacle_name != 'forbidden':   
                    row_base_height += height_offset_px
                    add_flag = False
                # goals[:, 0] += (obstacle_idx+1) * (self.track_block_resolution[0] * self.cfg.horizontal_scale)
                # self.goals[row_idx, col_idx, 1+(obstacle_idx)*2:1+(obstacle_idx+1)*2, :] = goals
            
            terrain.height_field_raw[
                heightfield_x0: heightfield_x1,
                heightfield_y0: heightfield_y1,
            ] = track_heightfield
            
            terrain.virtual_height_field[
                heightfield_x0: heightfield_x1,
                heightfield_y0: heightfield_y1,
            ] = virtual_height_field
            # self.track_info_map[row_idx, col_idx, local_row_id + 1, local_col_id, 0] = obstacle_id
            # self.track_info_map[row_idx, col_idx, local_row_id + 1, local_col_id, 1:] = block_info
            # self.next_block_height_map[row_idx, col_idx, local_row_id + 1, local_col_id] = block_starting_height_px * self.cfg.vertical_scale
            # block_starting_height_px += height_offset_px
    # if final_goal is not None:
    #     for item in goal_lists:
    #         # item = (item[0]+1, item[1])
    #         terrain_goals = np.vstack((terrain_goals, temp_goals[item]))
    terrain.goals = terrain_goals[1:]
            
def convert_heightfield_to_trimesh_delatin(height_field_raw, horizontal_scale, vertical_scale, max_error=0.01):
    # mesh = Delatin(np.flip(height_field_raw, axis=1).T, z_scale=vertical_scale, max_error=max_error)
    # vertices = np.zeros_like(mesh.vertices)
    # vertices[:, :2] = mesh.vertices[:, :2] * horizontal_scale
    # vertices[:, 2] = mesh.vertices[:, 2]
    # return vertices, mesh.triangles
    pass

def convert_heightfield_to_trimesh(height_field_raw, horizontal_scale, vertical_scale, slope_threshold=None):
    """
    Convert a heightfield array to a triangle mesh represented by vertices and triangles.
    Optionally, corrects vertical surfaces above the provide slope threshold:

        If (y2-y1)/(x2-x1) > slope_threshold -> Move A to A' (set x1 = x2). Do this for all directions.
                   B(x2,y2)
                  /|
                 / |
                /  |
        (x1,y1)A---A'(x2',y1)

    Parameters:
        height_field_raw (np.array): input heightfield
        horizontal_scale (float): horizontal scale of the heightfield [meters]
        vertical_scale (float): vertical scale of the heightfield [meters]
        slope_threshold (float): the slope threshold above which surfaces are made vertical. If None no correction is applied (default: None)
    Returns:
        vertices (np.array(float)): array of shape (num_vertices, 3). Each row represents the location of each vertex [meters]
        triangles (np.array(int)): array of shape (num_triangles, 3). Each row represents the indices of the 3 vertices connected by this triangle.
    """
    hf = height_field_raw
    num_rows = hf.shape[0]
    num_cols = hf.shape[1]

    y = np.linspace(0, (num_cols-1)*horizontal_scale, num_cols)
    x = np.linspace(0, (num_rows-1)*horizontal_scale, num_rows)
    yy, xx = np.meshgrid(y, x)

    if slope_threshold is not None:
        
        slope_threshold *= horizontal_scale / vertical_scale
        move_x = np.zeros((num_rows, num_cols))
        move_y = np.zeros((num_rows, num_cols))
        move_corners = np.zeros((num_rows, num_cols))
        move_x[:num_rows-1, :] += (hf[1:num_rows, :] - hf[:num_rows-1, :] > slope_threshold)
        move_x[1:num_rows, :] -= (hf[:num_rows-1, :] - hf[1:num_rows, :] > slope_threshold)
        move_y[:, :num_cols-1] += (hf[:, 1:num_cols] - hf[:, :num_cols-1] > slope_threshold)
        move_y[:, 1:num_cols] -= (hf[:, :num_cols-1] - hf[:, 1:num_cols] > slope_threshold)
        move_corners[:num_rows-1, :num_cols-1] += (hf[1:num_rows, 1:num_cols] - hf[:num_rows-1, :num_cols-1] > slope_threshold)
        move_corners[1:num_rows, 1:num_cols] -= (hf[:num_rows-1, :num_cols-1] - hf[1:num_rows, 1:num_cols] > slope_threshold)
        xx += (move_x + move_corners*(move_x == 0)) * horizontal_scale
        yy += (move_y + move_corners*(move_y == 0)) * horizontal_scale
    
    # if slope_threshold[1] is not None and slope_threshold[2] is not None:
    #     forbidden_edge = generate_forbidden_area(horizontal_scale, vertical_scale, num_rows, num_cols, hf, slope_threshold[1])
    #     obstacle_area = generate_forbidden_area(horizontal_scale, vertical_scale, num_rows, num_cols, hf, slope_threshold[2])
    #     obstacle_area = fill_rectangles(obstacle_area)
    #     # forbidden_area = forbidden_area_1 | forbidden_area_2
        

    # create triangle mesh vertices and triangles from the heightfield grid
    vertices = np.zeros((num_rows*num_cols, 3), dtype=np.float32)
    vertices[:, 0] = xx.flatten()
    vertices[:, 1] = yy.flatten()
    vertices[:, 2] = hf.flatten() * vertical_scale
    triangles = -np.ones((2*(num_rows-1)*(num_cols-1), 3), dtype=np.uint32)
    for i in range(num_rows - 1):
        ind0 = np.arange(0, num_cols-1) + i*num_cols
        ind1 = ind0 + 1
        ind2 = ind0 + num_cols
        ind3 = ind2 + 1
        start = 2*i*(num_cols-1)
        stop = start + 2*(num_cols-1)
        triangles[start:stop:2, 0] = ind0
        triangles[start:stop:2, 1] = ind3
        triangles[start:stop:2, 2] = ind1
        triangles[start+1:stop:2, 0] = ind0
        triangles[start+1:stop:2, 1] = ind2
        triangles[start+1:stop:2, 2] = ind3
    edge_mask = move_x != 0
    edge_mask |= move_y != 0
    return vertices, triangles, edge_mask